import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from matplotlib.patches import Polygon
from matplotlib.colors import rgb2hex
import numpy as np
import pandas as pd

plt.figure(figsize=(16, 8))
m = Basemap(
    llcrnrlon=77,
    llcrnrlat=14,
    urcrnrlon=140,
    urcrnrlat=51,
    projection='lcc',
    lat_1=33,
    lat_2=45,
    lon_0=100
)
m.drawcountries(linewidth=1.5)
m.drawcoastlines()

m.readshapefile('CHN_adm_shp/CHN_adm1', 'states', drawbounds=True)

df = pd.read_csv('pop.csv')
df['省名'] = df.地区
df.set_index('省名', inplace=True)

colors = {}
provinces = []
for shape_dict in m.states_info:
    name = shape_dict['NL_NAME_1']
    p = name.split('|')
    s = p[1] if len(p) > 1 else p[0]
    provinces.append(s)
    pop = df['人口数'][s]
    colors[s] = plt.cm.YlOrRd(np.sqrt((pop - 3000000) / (100000000 - 3000000)))[:3]

ax = plt.gca()
for n_shape, seg in enumerate(m.states):
    color = rgb2hex(colors[provinces[n_shape]])
    poly = Polygon(seg, facecolor=color, edgecolor=color)
    ax.add_patch(poly)

if __name__ == '__main__':
    plt.show()
